#!/usr/bin/env python3
"""修复 diffusers attention_dispatch.py 的 enable_gqa 参数问题"""
import os
import sys

try:
    import diffusers
    target = os.path.join(os.path.dirname(diffusers.__file__), 
                         'models/attention_dispatch.py')
    
    if not os.path.exists(target):
        print(f"File not found: {target}")
        sys.exit(0)
    
    print(f"Reading {target}...")
    with open(target, 'r', encoding='utf-8') as f:
        content = f.read()
    
    if '_npu_attention_patched' in content:
        print("Already patched")
        sys.exit(0)
    
    # 查找并替换 enable_gqa 参数
    if 'enable_gqa=' in content:
        print("Found enable_gqa parameter, removing it...")
        
        # 替换所有 enable_gqa 参数
        import re
        # 匹配 enable_gqa=xxx, 或 enable_gqa=xxx)
        content = re.sub(r',\s*enable_gqa\s*=\s*[^,\)]+', '', content)
        content = re.sub(r'enable_gqa\s*=\s*[^,\)]+,\s*', '', content)
        
        # 添加标记
        content = '# NPU compatibility patch - _npu_attention_patched\n' + content
        
        with open(target, 'w', encoding='utf-8') as f:
            f.write(content)
        
        print(f"Successfully patched {target}")
        
        # 验证
        with open(target, 'r', encoding='utf-8') as f:
            verify = f.read()
        
        if 'enable_gqa=' in verify:
            print("WARNING: enable_gqa still exists after patch")
            sys.exit(1)
        
        print("Patch verified successfully")
    else:
        print("enable_gqa parameter not found, no patch needed")
        
except Exception as e:
    import traceback
    print(f"Error: {e}", file=sys.stderr)
    traceback.print_exc(file=sys.stderr)
    sys.exit(1)
